package com.sy.nlp

import com.hankcs.hanlp.HanLP
import com.hankcs.hanlp.seg.common.Term
import org.apache.spark.sql.SparkSession

import java.util
import java.util.StringJoiner
import scala.collection.JavaConversions._

object HanlpTest {
  case class CalculateResult(text: String, time: String, person: String, byPerson: String, location: String, action: String, keyWord: String, summary: String)

  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .appName("test")
      .master("local[*]")
      .config("spark.driver.memory", "12G")
      .config("spark.kryoserializer.buffer.max", "200M")
      .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
      .getOrCreate()

    spark.sparkContext.setLogLevel("WARN")

    val names = NameUtils.names()

    import spark.implicits._


    spark.read.option("encoding", "gbk").csv("data.csv").limit(1000)
      .selectExpr("_c3 as text", "_c1 as time", "_c0 as zhudong_person")
      .rdd
      .map(sentence => {
        val segment = HanLP.newSegment
          .enableNameRecognize(true)
          .enablePlaceRecognize(true)
          .enableOrganizationRecognize(true)
        val text = sentence.getString(0)
        val time = sentence.getString(1)
        val person = sentence.getString(2)

        val personJsonObj = new StringJoiner("。")
        val placeJsonObj = new StringJoiner("。")
        val actionJsonObj = new StringJoiner("。")
        val keyWordJsonObj = new StringJoiner("。")
        val summaryJsonObj = new StringJoiner("。")

        if (text != null) {

          var index = 1
          for (item <- text.split("。")) {
            val personList = new util.ArrayList[String]()
            val placeList = new util.ArrayList[String]()
            val actionList = new util.ArrayList[String]()

            val termList: util.List[Term] = segment.seg(item)

            for (term <- termList) {
              if (Constant.NLP_PERSON_MAP.containsKey(term.nature) && !names.contains(term.word)) {
                personList.add(term.word)
              }
              if (Constant.NLP_PLACE_MAP.containsKey(term.nature)) {
                placeList.add(term.word)
              }
              if (Constant.NLP_ACTION_MAP.containsKey(term.nature)) {
                actionList.add(term.word)
              }

            }


            val keyWord = HanLP.extractKeyword(item, 3)
            val summary = HanLP.extractSummary(item, 10)

            personJsonObj.add(index + "、" + personList.toString)
            placeJsonObj.add(index + "、" + placeList)
            actionJsonObj.add(index + "、" + actionList)
            keyWordJsonObj.add(index + "、" + keyWord)
            summaryJsonObj.add(index + "、" + summary)

            index = index + 1

          }
        }
        CalculateResult(
          text,
          time,
          person,
          personJsonObj.toString,
          placeJsonObj.toString,
          actionJsonObj.toString,
          keyWordJsonObj.toString,
          summaryJsonObj.toString
        )
      }).toDF()
      .write
      .csv("/Users/wangsixian/code/result")
  }

}
